import cv2
import random
import torchvision.transforms.functional as FT


class ColorConvert(object):
    """
    input image is np.ndarray, not PIL
    """
    def __init__(self, src_mode, dst_mode):
        self.src_mode = src_mode
        self.dst_mode = dst_mode
    
    def __call__(self, image, boxes=None, labels=None):
        if self.src_mode == 'RGB' and self.dst_mode == 'HSV':
            image = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)
        elif self.src_mode == 'RGB' and self.dst_mode == 'BGR':
            image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
        elif self.src_mode == 'BGR' and self.dst_mode == 'HSV':
            image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
        elif self.src_mode == 'BGR' and self.dst_mode == 'RGB':
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        elif self.src_mode == 'HSV' and self.dst_mode == 'RGB':
            image = cv2.cvtColor(image, cv2.COLOR_HSV2RGB)
        elif self.src_mode == 'HSV' and self.dst_mode == 'BGR':
            image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR)
        else:
            raise NotImplementedError

        return image, boxes, labels


class RandomChannelPermutation(object):
    """
    input image is np.ndarray, not PIL
    """
    def __init__(self):
        self.perms = ((0, 1, 2), (0, 2, 1), 
                      (1, 0, 2), (1, 2, 0),
                      (2, 0, 1), (2, 1, 0))
    
    def __call__(self, image, boxes=None, labels=None):
        if random.randint(2):
            swap_indices = self.perms[random.randint(6)]
            image = image[:, :, swap_indices]
        return image, boxes, labels

# -----------------------------------------------------------------------------
class RandomBrightness(object):
    def __init__(self, min_delta=0.3, max_delta=3.0):
        self.min_delta = min_delta
        self.max_delta = max_delta

    def __call__(self, image, boxes=None, labels=None):
        delta = random.uniform(self.min_delta, self.max_delta)
        image = FT.adjust_brightness(image, delta)
        return image, boxes, labels


class RandomContrast(object):
    def __init__(self, min_delta=0.3, max_delta=3.0):
        self.min_delta = min_delta
        self.max_delta = max_delta
    
    def __call__(self, image, boxes=None, labels=None):
        delta = random.uniform(self.min_delta, self.max_delta)
        image = FT.adjust_contrast(image, delta)

        return image, boxes, labels


class RandomSaturation(object):
    def __init__(self, min_delta=0.3, max_delta=3.0):
        self.min_delta = min_delta
        self.max_delta = max_delta
    
    def __call__(self, image, boxes=None, labels=None):
        delta = random.uniform(self.min_delta, self.max_delta)
        image = FT.adjust_saturation(image, delta)

        return image, boxes, labels


class RandomHue(object):
    def __init__(self, max_delta=18./255):
        self.max_delta = max_delta

    def __call__(self, image, boxes=None, labels=None):
        delta = random.uniform(-self.max_delta, self.max_delta)
        image = FT.adjust_hue(image, delta)

        return image, boxes, labels


class ToGrayScale(object):
    def __call__(self, image, boxes, labels):
        pass


class RandomChannelPermute(object):
    def __init__(self):
        pass

    def __call__(self, image, boxes, labels):
        pass


class RandomPhotometric(object):
    """
    apply random photometric distortions.
    """
    def __init__(self, p=0.5):
        self.prob = p
        self.distortions = [RandomBrightness(), RandomHue(),
                            RandomSaturation(), RandomContrast()]
    
    def __call__(self, image, boxes, labels):
        for d in self.distortions:
            if random.random() < self.prob:
                image, boxes, labels = d(image, boxes, labels)
        return image, boxes, labels
